package org.example.aitest.test;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.tools.FunctionDefinition;
import com.alibaba.dashscope.tools.ToolCallBase;
import com.alibaba.dashscope.tools.ToolCallFunction;
import com.alibaba.dashscope.tools.ToolFunction;
import com.alibaba.dashscope.utils.JsonUtils;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.github.victools.jsonschema.generator.*;
import io.reactivex.Flowable;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;

public class Main {

    public static Flowable<GenerationResult> callWithMessage() throws NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("百炼什么")
                .build();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content("今天是几号")
                .build();
        Message systemMsg1 = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("你好，Peeter！很高兴认识你。如果你有任何问题或需要帮助的地方，请随时告诉我哦！我是一个超大规模的语言模型，可以帮你解答各种问题、创作文字，比如写故事、写公文、写邮件、写剧本，进行逻辑推理、编程等等。有什么想聊的吗？")
                .build();
        Message userMsg1 = Message.builder()
                .role(Role.USER.getValue())
                .content("我是谁")
                .build();
        Message systemMsg2 = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("你好，Peeter！你是独特的个体，可以是我交谈的朋友，提出问题的探索者，或是创作内容的合作者。我是一个超大规模的语言模型，我的职责是帮助你，无论你是谁，无论你需要什么样的帮助——解答疑惑、创作文字、逻辑推理、编程等都在我的能力范围内。那么，Peeter，你今天想做些什么呢？")
                .build();
        Message userMsg2 = Message.builder()
                .role(Role.USER.getValue())
                .content("直接回答我的问题就行，不要说其他的，现在回答我是谁")
                .build();
        Message systemMsg3 = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("你是Peeter，一个正在与我交谈的朋友。")
                .build();
        Message userMsg3 = Message.builder()
                .role(Role.USER.getValue())
                .content("只回答我问的,不要有多余语言")
                .build();

        GenerationParam param = GenerationParam.builder()
                .apiKey("sk-06daae4607c0459295ccd4db313f9719")
                .model("qwen-plus")
//                .incrementalOutput(true) //流式输出时是否以增量输出，默认为false
                .messages(Arrays.asList(systemMsg,userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();

//        gen.call(param); //非流式输出
        return gen.streamCall(param);
    }

//    public static void main(String[] args) {
//        try {
//            Flowable<GenerationResult> resultFlowable = callWithMessage();
//
////            // 流式输出
////            resultFlowable.blockingForEach(message -> {
////                String content = message.getOutput().getChoices().get(0).getMessage().getContent();
////                System.out.println(content);
////            });
//            System.out.println(resultFlowable.blockingLast().getOutput().getChoices().get(0).getMessage().getContent());
////            System.out.println(result.getOutput().getChoices().get(0).getMessage().getContent());
//        }catch (ApiException | NoApiKeyException | InputRequiredException e) {
//            System.err.println("错误信息："+e.getMessage());
//            System.out.println("请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code");
//        }
//        System.exit(0);
//    }

    public static void main(String[] args) {
        try {
            selectTool();
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            System.out.println(String.format("Exception: %s", e.getMessage()));
        } catch (Exception e) {
            System.out.println(String.format("Exception: %s", e.getMessage()));
        }
        System.exit(0);
    }

    public static class GetWeatherTool {
        private String location;

        public GetWeatherTool(String location) {
            this.location = location;
        }

        public String callWeather() {
            // 假设location是一个JSON字符串，例如{"location": "北京"}
            // 需要提取其中的"location"的值
            location = "{\"location\": \"北京\"}";
            try {
                // 使用Jackson库解析JSON
                ObjectMapper objectMapper = new ObjectMapper();
                JsonNode jsonNode = objectMapper.readTree(location);
                String locationName = jsonNode.get("location").asText();
                return locationName + "今天是晴天";
            } catch (Exception e) {
                // 如果解析失败，返回原始字符串
                return location + "今天是晴天";
            }
        }
    }

    public static class GetTimeTool {
        public String getCurrentTime() {
            LocalDateTime now = LocalDateTime.now();
            DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
            return "当前时间：" + now.format(formatter) + "。";
        }
    }

    private static ObjectNode generateSchema(Class<?> clazz) {
        SchemaGeneratorConfigBuilder configBuilder =
                new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON);
        SchemaGeneratorConfig config = configBuilder.with(Option.EXTRA_OPEN_API_FORMAT_VALUES)
                .without(Option.FLATTENED_ENUMS_FROM_TOSTRING).build();
        SchemaGenerator generator = new SchemaGenerator(config);
        return generator.generateSchema(clazz);
    }

    public static void selectTool()
            throws NoApiKeyException, ApiException, InputRequiredException {
        ObjectNode jsonSchemaWeather = generateSchema(GetWeatherTool.class);
        ObjectNode jsonSchemaTime = generateSchema(GetTimeTool.class);
        FunctionDefinition fdWeather = FunctionDefinition.builder().name("get_current_weather")
                .description("获取指定地区的天气")
                .parameters(JsonUtils.parseString(jsonSchemaWeather.toString()).getAsJsonObject()).build();
        FunctionDefinition fdTime = FunctionDefinition.builder().name("get_current_time")
                .description("获取当前时刻的时间")
                .parameters(JsonUtils.parseString(jsonSchemaTime.toString()).getAsJsonObject()).build();
        Message systemMsg = Message.builder().role(Role.SYSTEM.getValue())
                .content("You are a helpful assistant. When asked a question, use tools wherever possible.")
                .build();
        Scanner scanner = new Scanner(System.in);
        System.out.println("请输入：");
        String userInput = scanner.nextLine();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(userInput).build();
        List<Message> messages = new ArrayList<>(Arrays.asList(systemMsg, userMsg));
        GenerationParam param = GenerationParam.builder()
                // 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .model("qwen-plus")
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey("sk-06daae4607c0459295ccd4db313f9719")
                .messages(messages).resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .tools(Arrays.asList(
                        ToolFunction.builder().function(fdWeather).build(),
                        ToolFunction.builder().function(fdTime).build()
                )).build();
        Generation gen = new Generation();
        GenerationResult result = gen.call(param);
        System.out.println("首轮输出：" + JsonUtils.toJson(result));
        boolean needToolCall = true;
        while (needToolCall) {
            needToolCall = false;
            for (GenerationOutput.Choice choice : result.getOutput().getChoices()) {
                messages.add(choice.getMessage());
                if (choice.getMessage().getToolCalls() != null) {
                    for (ToolCallBase toolCall : choice.getMessage().getToolCalls()) {
                        if (toolCall.getType().equals("function")) {
                            String functionName = ((ToolCallFunction) toolCall).getFunction().getName();
                            String functionArgument = ((ToolCallFunction) toolCall).getFunction().getArguments();
                            if (functionName.equals("get_current_weather")) {
                                GetWeatherTool weatherTool = new GetWeatherTool(functionArgument);
                                String weather = weatherTool.callWeather();
                                Message toolResultMessage = Message.builder().role("tool")
                                        .content(weather).toolCallId(toolCall.getId()).build();
                                messages.add(toolResultMessage);
                                System.out.println("工具输出信息：" + weather);
                            } else if (functionName.equals("get_current_time")) {
                                GetTimeTool timeTool = new GetTimeTool();
                                String time = timeTool.getCurrentTime();
                                Message toolResultMessage = Message.builder().role("tool")
                                        .content(time).toolCallId(toolCall.getId()).build();
                                messages.add(toolResultMessage);
                                System.out.println("工具输出信息：" + time);
                            }
                            needToolCall = true;
                        }
                    }
                } else {
                    System.out.println("最终答案：" + choice.getMessage().getContent());
                    return;
                }
            }

            if (needToolCall) {
                param.setMessages(messages);
                result = gen.call(param);
                System.out.println("下一轮输出：" + JsonUtils.toJson(result));
            }
        }

        System.out.println("最终答案：" + result.getOutput().getChoices().get(0).getMessage().getContent());
    }

}
